/**
 * 
 */
package edu.umd.clip.lm.model.training;

import edu.berkeley.nlp.util.*;
import edu.umd.clip.lm.model.training.NewTrainer.Distributions;
import edu.umd.clip.lm.util.*;

/**
 * @author Denis Filimonov <den@cs.umd.edu>
 *
 */
public class CrossEntropyQuestionEstimator extends AbstractQuestionEstimator {
	private final int minCount;
	private final double minCostReduction;
	private static final int MAX_NR_ITERATIONS = 200;
	
	public CrossEntropyQuestionEstimator(int minCount, double minCostReduction) {
		this.minCount = minCount;
		this.minCostReduction = minCostReduction;
	}

	private static double getInitialLambda(ActiveTreeNode activeNode) {
		double initialLambda;
		
		initialLambda = activeNode.lambda > 0 ? Math.sqrt(activeNode.lambda / 2) : 0.9;
		if (initialLambda > 0.95) return 0.95;
		if (initialLambda < 0.05) return 0.05;
		return initialLambda;
	}

	static double computeCrossEntropy(Long2DoubleMap parentDist, 
			Long2DoubleMap trainingDist, long trainingCount, 
			Long2DoubleMap devDist, long devCount,
			MutableDouble Lambda, double initialLambda) 
	{
		
		// filter out irrelevant stuff
		/*
		Long2DoubleMap trimmedParentDist = new Long2DoubleMap(devDist.size());
		Long2DoubleMap trimmedTrainingDist = new Long2DoubleMap(devDist.size());
		
		for(Long2DoubleMap.Iterator i = devDist.iterator(); i.hasNext();) {
			Long2DoubleMap.Entry entry = i.next();
			long key = entry.getKey();
			trimmedParentDist.put(key, parentDist.get(key));
			trimmedTrainingDist.put(key, trainingDist.get(key));
		}
		*/
		//trainingDist = null;
		
		double revTrainingCount = 1.0 / trainingCount;
		double revDevCount = 1.0 / devCount;
		// do the EM
		double lambda = initialLambda;
		int iterations = 0;
		while(iterations < MAX_NR_ITERATIONS) {
			++iterations;
			double sum = 0;
			for(Long2DoubleMap.Iterator i = devDist.iterator(); i.hasNext();) {
				Long2DoubleMap.Entry entry = i.next();
				long key = entry.getKey();
				double p = trainingDist.get(key) * revTrainingCount;
				double parent_p = parentDist.get(key);
				double interp_p = lambda * p + (1-lambda) * parent_p;
				
				double add = entry.getValue() * p / interp_p;
				if (Double.isNaN(add)) {
					//System.err.printf("tuple %s not found in training distribution.\n", FactorTuple.toString(key));
					continue;
				}
				sum += add;
			}
			sum *= revDevCount;
			
			double new_lambda = lambda * sum;
			lambda = new_lambda;
			
			if (Math.abs(sum - 1.0) < 0.01) {
				break;
			}
			if (lambda < 1e-6) {
				break;
			}
		}
		Lambda.set(lambda);
		
		// compute cross entropy
		double entropy = 0;
		for(Long2DoubleMap.Iterator i = devDist.iterator(); i.hasNext();) {
			Long2DoubleMap.Entry entry = i.next();
			long key = entry.getKey();
			double p = trainingDist.get(key);
			double parent_p = parentDist.get(key);
			double interp_p = lambda * p + (1-lambda) * parent_p;

			if (interp_p > 0) {
				entropy -= entry.getValue() * ProbMath.log2(interp_p);
			}
		}
		
		//System.out.printf("EM: %d iterations, lambda = %g (initial %g), entropy = %g\n", iterations, lambda, initialLambda, entropy);
		return entropy;
	}

	public static class CrossEntropyResult extends Result {
		double leftLambda;
		double rightLambda;
		/**
		 * @param cost
		 */
		public CrossEntropyResult(double cost) {
			super(cost);
		}
		public double getLeftLambda() {
			return leftLambda;
		}
		public double getRightLambda() {
			return rightLambda;
		}
		@Override
		public void setupFalseActiveNode(int dataIdx, ActiveTreeNode activeNode, ActiveTreeNode parentNode, Distributions distributions) {
			throw new Error("not implemented");
			/*
			super.setupFalseActiveNode(dataIdx, activeNode, parentNode, distributions);
			activeNode.lambda = leftLambda;
			activeNode.smoothedDistributions[dataIdx] = getTrimmedDistribution(parentNode.smoothedDistributions[dataIdx], 
					distributions.falseDist, distributions.falseCount, distributions.falseDev, leftLambda);
			*/
		}
		@Override
		public void setupTrueActiveNode(int dataIdx, ActiveTreeNode activeNode, ActiveTreeNode parentNode, Distributions distributions) {
			throw new Error("not implemented");
			/*
			super.setupTrueActiveNode(dataIdx, activeNode, parentNode, distributions);
			activeNode.lambda = rightLambda;
			activeNode.smoothedDistribution = getTrimmedDistribution(parentNode.smoothedDistribution, 
					distributions.trueTrain, distributions.trueTrainCount, distributions.trueDev, rightLambda);
			*/
		}
		
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.AbstractQuestionEstimator#estimateQuestion(edu.umd.clip.lm.model.training.ActiveTreeNode, edu.umd.clip.lm.model.training.NewTrainer.Distributions)
	 */
	@Override
	public Result estimateQuestion(int data, ActiveTreeNode activeNode, Distributions distributions) {
		throw new Error("not implemented");
		/*
		MutableDouble newTrueLambda = new MutableDouble();
		MutableDouble newFalseLambda = new MutableDouble();
		
		double initialLambda = getInitialLambda(activeNode);
		
		double trueEntropy = computeCrossEntropy(activeNode.smoothedDistribution,
				distributions.trueTrain, distributions.trueTrainCount,
				distributions.trueDev, distributions.trueDevCount,
				newTrueLambda, initialLambda);
		
		double falseEntropy = computeCrossEntropy(activeNode.smoothedDistribution,
				distributions.falseTrain, distributions.falseTrainCount,
				distributions.falseDev, distributions.falseDevCount,
				newFalseLambda, initialLambda);
		
		double entropy = (falseEntropy*distributions.falseTrainCount + trueEntropy*distributions.trueTrainCount) /
			(distributions.falseTrainCount + distributions.trueTrainCount);
		
		CrossEntropyResult result = new CrossEntropyResult(entropy);
		result.falseCost = falseEntropy;
		result.trueCost = trueEntropy;
		result.leftLambda = newFalseLambda.doubleValue();
		result.rightLambda = newTrueLambda.doubleValue();
		return result;
		*/
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.AbstractQuestionEstimator#isGood(edu.umd.clip.lm.model.training.AbstractQuestionEstimator.Result, edu.umd.clip.lm.model.training.ActiveTreeNode)
	 */
	@Override
	public boolean isGood(int data, Result result, ActiveTreeNode activeNode) {
		return activeNode.nodeCosts[data] - result.cost > minCostReduction;
	}

	/* (non-Javadoc)
	 * @see edu.umd.clip.lm.model.training.AbstractQuestionEstimator#prePruningRule(edu.umd.clip.lm.model.training.ActiveTreeNode)
	 */
	@Override
	public boolean prePruningRule(ActiveTreeNode activeNode) {
		for(long count :activeNode.counts) {
			if (count < minCount) return false;
		}
		return true;
	}

	private static Long2DoubleMap getTrimmedDistribution(Long2DoubleMap parentDist, 
			Long2DoubleMap childDist, long childDistCount, 
			Long2DoubleMap devDist,	double lambda) 
	{
		Long2DoubleMap trimmedDist = new Long2DoubleMap(devDist.size());

		double revChildDistCount = 1.0 / childDistCount;
		
		for(Long2DoubleMap.Iterator i = devDist.iterator(); i.hasNext();) {
			Long2DoubleMap.Entry entry = i.next();
			long key = entry.getKey();
			double p = lambda * childDist.get(key) * revChildDistCount + (1-lambda) * parentDist.get(key);
			trimmedDist.put(key, p);
		}
		return trimmedDist;
	}
	
}
